package com.example.smartgreenhouse.aop;

import com.example.smartgreenhouse.annotation.Limit;
import com.example.smartgreenhouse.utils.JwtUtil;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.reflect.MethodSignature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletResponse;
import java.io.PrintWriter;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.Objects;

/**
 * 使用AOP切面拦截限流注解
 */
@Slf4j
@Aspect
@Component
public class LimitAop {

    @Autowired
    private JwtUtil jwtUtil;
    /**
     * 不同的接口，不同的流量控制
     * map的key为 Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    //使用 @Around 注解来定义环绕通知，针对带有 @Limit 注解的方法进行拦截。
    @Around("@annotation(com.example.smartgreenhouse.annotation.Limit)")  //表示要拦截带有某个注解的方法
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable{

        /* 这是从切入点获取方法签名（Signature 类型），但是：
        它只是一个通用接口，不能直接获取方法对象（Method）
        (MethodSignature) 强转成 MethodSignature
        MethodSignature 是 Signature 的子类，代表一个具体的方法签名。*/
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        //获取方法上的 Limit 注解，以获取相关的限流配置。
        Limit limit = method.getAnnotation(Limit.class);
        String key;
        if (limit != null) {
            //处理针对IP的限流，利用IP和limit中的key生成唯一的key。
            if( limit.limitType() == Limit.LimitType.IP){
                //获取请求中的IP。
                String ip = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes()))
                        .getRequest().getRemoteAddr();
                key = limit.key() + "_" + ip;
            }
            //处理针对User的限流，利用用户的ID和limit中的key生成唯一的key。
            else if (limit.limitType() == Limit.LimitType.USER){
                //获取用户的ID。
                String token = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder
                        .getRequestAttributes()))
                        .getRequest()
                        .getHeader("token");
                Map<String,String> map = jwtUtil.checkSign(token);
                key = limit.key()+"_"+map.get("userId");
            }
            //处理没有特定规定的限流，仅仅使用limit中的key
            else{
                key = limit.key();
            }
            // 使用 computeIfAbsent 保证原子性，创建和获取令牌桶
            RateLimiter rateLimiter = limitMap.computeIfAbsent(key, k -> {
                log.info("新建了令牌桶={}，容量={}", k, limit.permitsPerSecond());
                return RateLimiter.create(limit.permitsPerSecond());
            });
            //通过 tryAcquire 方法尝试获取令牌，如果失败则返回 429 响应。
            boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
            if (!acquire) {
                responseFail(limit.msg());
                return null;
            }
        }
        return joinPoint.proceed();
    }

    /**
     * 直接向前端抛出异常
     * @param msg 提示信息
     */
    //该方法用于获取令牌失败时的响应处理，其中返回状态码429,表示请求过于频繁
    private void responseFail(String msg) {
        HttpServletResponse response;
        try {
            response = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder
                    .getRequestAttributes()))
                    .getResponse();
            if (response != null) {
                response.setCharacterEncoding("UTF-8");
                response.setContentType("application/json");
                response.setStatus(429); // 设置状态码为 429
                // 使用 try-with-resources 自动关闭 PrintWriter
                try (PrintWriter pr = response.getWriter()) {
                    pr.write("{\"code\": \"429\",\"status\": \"error\", \"message\": \"" + msg + "\"}");
                    pr.flush();
                }
            }
        } catch (Exception e) {
            log.error("响应失败信息时出现异常", e);
        }
    }
}